package tcpsyn

import (
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"nprobe/core"
	"nprobe/utils/routines"
	"time"
)

type TCPSYNReceiver struct {
	core.BaseReceiver
	Bucket *core.PktRecordBucket
}

func NewTCPSYNReceiver(
	ctx context.Context, name string, id int, task *core.ProbeTask, salts map[int][]byte, targetSet map[string]bool, bucket *core.PktRecordBucket,
) *TCPSYNReceiver {
	ppRBuf := core.NewPPRBuf(task.SpanMs, task.TimeoutMs)

	return &TCPSYNReceiver{
		BaseReceiver: core.BaseReceiver{
			Ctx:      ctx,
			Name:     name,
			ID:       id,
			Task:     task,
			RBuf:     ppRBuf,
			Salts:    salts,
			DstIPSet: targetSet,
		},
		Bucket: bucket,
	}
}

func (r *TCPSYNReceiver) StartReceiver() error {
	routines.GoWithRecover(r.Name, r.doReceiveTCP)
	return nil
}

func (r *TCPSYNReceiver) doReceiveTCP() {
	golog.Debugf("UDPReceiver doReceiveTCP")

	conn, err := core.OpenReceiverIP4PacketConn("ip4:tcp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}

	r.Conn = conn

	filter := core.CreateBPFTCPSYN(r.Task.DstPortMin, r.Task.DstPortMax, r.Task.SrcPortMin, r.Task.SrcPortMax)
	err = core.SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := core.AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			conn.Close()
			return
		default:
			r.receiveTCPPkt(payloadBuf)
		}
	}
}

func (r *TCPSYNReceiver) receiveTCPPkt(payloadBuf []byte) {
	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		golog.Debugf("err: %+v", err)
		return
	}
	rxTs := time.Now().UnixNano()

	// ip header
	srcIP := h.Src
	dstIP := h.Dst
	ipID := uint16(h.ID)
	devIP := srcIP.To4().String()

	// tcp header
	srcPort := binary.BigEndian.Uint16(payload[0:2])
	dstPort := binary.BigEndian.Uint16(payload[2:4])
	tcpACK := binary.BigEndian.Uint32(payload[8:12])
	tcpFlags := payload[13]

	if (tcpFlags & 0x04) != 0 {
		// RST
		ipID = uint16(tcpACK - uint32(r.Task.MsgLen) - 1)
	} else if (tcpFlags & 0x012) == 0x12 {
		// SYN+ACK
		ipID = uint16(tcpACK - 1)
	} else if (tcpFlags & 0x010) == 0x10 {
		// ACK
		ipID = uint16(tcpACK - 1)
	} else {
		return
	}

	if !r.DstIPSet[devIP] {
		return
	}

	r.recordPkt(dstIP.To4().String(), devIP, dstPort, srcPort, ipID, rxTs, false)
}

func (r *TCPSYNReceiver) recordPkt(localIP string, remoteIP string, srcPort uint16, dstPort uint16, ipID uint16, rxTs int64, bitFlip bool) {
	golog.Debugf("%s: %s:%d->%s:%d, ipID: %d, txTs: %d", r.Name, localIP, srcPort, remoteIP, dstPort, ipID, rxTs)
	key := fmt.Sprintf("%s-%d-%d-%d", remoteIP, ipID, srcPort, dstPort)
	pktRecord := &core.PktRecord{
		Key:      key,
		SrcIP:    localIP,
		DstIP:    remoteIP,
		Proto:    r.Task.Proto,
		Tos:      r.Task.TOS,
		RxTs:     rxTs,
		Received: true,
		BitFlip:  bitFlip,
	}

	r.Bucket.Update(pktRecord)
}
